# -*- coding：utf-8 -*-
import base64
import random

import rsa
from Crypto.Cipher import AES
from Crypto.PublicKey import RSA
from pyDes import des, CBC, PAD_PKCS5
from Crypto.Cipher import DES3
import hashlib
import hmac
import sys
import time


class MersenneTwister:
    __n = 624
    __m = 397
    __a = 0x9908b0df
    __b = 0x9d2c5680
    __c = 0xefc60000
    __kInitOperand = 0x6c078965
    __kMaxBits = 0xffffffff
    __kUpperBits = 0x80000000
    __kLowerBits = 0x7fffffff

    def __init__(self, seed=0):
        self.__register = [0] * self.__n
        self.__state = 0

        self.__register[0] = seed
        for i in range(1, self.__n):
            prev = self.__register[i - 1]
            temp = self.__kInitOperand * (prev ^ (prev >> 30)) + i
            self.__register[i] = temp & self.__kMaxBits

    def __twister(self):
        for i in range(self.__n):
            y = (self.__register[i] & self.__kUpperBits) + \
                (self.__register[(i + 1) % self.__n] & self.__kLowerBits)
            self.__register[i] = self.__register[(i + self.__m) % self.__n] ^ (y >> 1)
            if y % 2:
                self.__register[i] ^= self.__a
        return None

    def __temper(self):
        if self.__state == 0:
            self.__twister()

        y = self.__register[self.__state]
        y = y ^ (y >> 11)
        y = y ^ (y << 7) & self.__b
        y = y ^ (y << 15) & self.__c
        y = y ^ (y >> 18)

        self.__state = (self.__state + 1) % self.__n

        return y

    def __call__(self):
        return self.__temper()

    def load_register(self, register):
        self.__state = 0
        self.__register = register


class switch(object):
    def __init__(self, value):
        self.value = value
        self.fall = False

    def __iter__(self):
        """Return the match method once, then stop"""
        try:
            yield self.match
        except StopIteration:
            return

    def match(self, *args):
        """Indicate whether or not to enter a case suite"""
        if self.fall or not args:
            return True
        elif self.value in args:  # changed for v1.5, see below
            self.fall = True
            return True
        else:
            return False


class USE_AES:
    """
    AES
    除了MODE_SIV模式key长度为：32, 48, or 64,
    其余key长度为16, 24 or 32
    详细见AES内部文档
    CBC模式传入iv参数
    本例使用常用的ECB模式
    """

    def __init__(self, key):
        if len(key) > 32:
            key = key[:32]
        self.key = self.to_16(key)

    def to_16(self, key):
        """
        转为16倍数的bytes数据
        :param key:
        :return:
        """
        key = bytes(key, encoding="utf8")
        while len(key) % 16 != 0:
            key += b'\0'
        return key  # 返回bytes

    def aes(self):
        return AES.new(self.key, AES.MODE_ECB)  # 初始化加密器

    def encrypt(self, text):
        aes = self.aes()
        return str(base64.encodebytes(aes.encrypt(self.to_16(text))),
                   encoding='utf8').replace('\n', '')  # 加密

    def decodebytes(self, text):
        aes = self.aes()
        return str(aes.decrypt(base64.decodebytes(bytes(
            text, encoding='utf8'))).rstrip(b'\0').decode("utf8"))  # 解密


class USE_RSA:
    """
    生成密钥可保存.pem格式文件
    1024位的证书，加密时最大支持117个字节，解密时为128；
    2048位的证书，加密时最大支持245个字节，解密时为256。
    加密大文件时需要先用AES或者DES加密，再用RSA加密密钥，详细见文档
    文档:https://stuvel.eu/files/python-rsa-doc/usage.html#generating-keys
    """

    def __init__(self, number=1024):
        """
        :param number: 公钥、私钥
        """
        self.pubkey, self.privkey = rsa.newkeys(number)

    def rsaEncrypt(self, text):
        """
        :param test: str
        :return: bytes
        """
        content = text.encode('utf-8')
        crypto = rsa.encrypt(content, self.pubkey)
        return crypto

    def rsaDecrypt(self, text):
        """
        :param text:bytes 
        :return: str
        """
        content = rsa.decrypt(text, self.privkey)
        con = content.decode('utf-8')
        return con

    def savePem(self, path_name, text):
        """
        :param path_name: 保存路径
        :param text: str
        :return:bytes
        """
        if "PEM" in path_name.upper():
            path_name = path_name[:-4]
        with open('{}.pem'.format(path_name), 'bw') as f:
            f.write(text.save_pkcs1())

    def readPem(self, path_name, key_type):
        """
        :param path_name: 密钥文件
        :param key_type:类型 
        :return: 
        """
        if 'pubkey' in key_type:
            self.pubkey = rsa.PublicKey.load_pkcs1(path_name)
        else:
            self.privkey = rsa.PublicKey.load_pkcs1(path_name)
        return True

    def sign(self, message, priv_key=None, hash_method='SHA-1'):
        """
        生成明文的哈希签名以便还原后对照
        :param message: str
        :param priv_key:
        :param hash_method: 哈希的模式
        :return:
        """
        if None == priv_key:
            priv_key = self.privkey
        return rsa.sign(message.encode(), priv_key, hash_method)

    def checkSign(self, mess, result, pubkey=None):
        """
        验证签名：传入解密后明文、签名、公钥，验证成功返回哈希方法，失败则报错
        :param mess: str
        :param result: bytes
        :param pubkey: 
        :return: str
        """
        if None == pubkey:
            pubkey = self.privkey
        try:
            result = rsa.verify(mess, result, pubkey)
            return result
        except:
            return False


class USE_DES:
    """
    des(key,[mode], [IV], [pad], [pad mode])
    key:必须正好8字节
    mode（模式）：ECB、CBC
    iv:CBC模式中必须提供长8字节
    pad:填充字符
    padmode:加密填充模式PAD_NORMAL or PAD_PKCS5
    """

    def __init__(self, key, iv):
        if not isinstance(key, bytes):
            key = bytes(key, encoding="utf8")
        if not isinstance(iv, bytes):
            iv = bytes(iv, encoding="utf8")
        self.key = key
        self.iv = iv

    def encrypt(self, text):
        """
        DES 加密
        :param text: 原始字符串
        :return: 加密后字符串，bytes
        """
        if not isinstance(text, bytes):
            text = bytes(text, "utf-8")
        secret_key = self.key
        iv = self.iv
        k = des(secret_key, CBC, iv, pad=None, padmode=PAD_PKCS5)
        en = k.encrypt(text, padmode=PAD_PKCS5)
        return en

    def descrypt(self, text):
        """
        DES 解密
        :param text: 加密后的字符串，bytes
        :return:  解密后的字符串
        """
        secret_key = self.key
        iv = self.iv
        k = des(secret_key, CBC, iv, pad=None, padmode=PAD_PKCS5)
        de = k.decrypt(text, padmode=PAD_PKCS5)
        return de.decode()


class USE_DES3:
    """
    new(key, mode, *args, **kwargs)
    key:必须8bytes倍数介于16-24
    mode：
    iv:初始化向量适用于MODE_CBC、MODE_CFB、MODE_OFB、MODE_OPENPGP，4种模式
        ``MODE_CBC``, ``MODE_CFB``, and ``MODE_OFB``长度为8bytes
        ```MODE_OPENPGP```加密时8bytes解密时10bytes
        未提供默认随机生成
    nonce：仅在 ``MODE_EAX`` and ``MODE_CTR``模式中使用
            ``MODE_EAX``建议16bytes
            ``MODE_CTR``建议[0, 7]长度
            未提供则随机生成
    segment_size：分段大小，仅在 ``MODE_CFB``模式中使用，长度为8倍数，未指定则默认为8
    mac_len： 适用``MODE_EAX``模式，身份验证标记的长度（字节），它不能超过8（默认值）
    initial_value：适用```MODE_CTR```，计数器的初始值计数器块。默认为**0**。
    """

    def __init__(self, key):
        self.key = key
        self.mode = DES3.MODE_ECB

    def encrypt(self, text):
        """
        传入明文
        :param text:bytes类型，长度是KEY的倍数
        :return:
        """
        if not isinstance(text, bytes):
            text = bytes(text, 'utf-8')
        x = len(text) % 8
        text = text + b'\0' * x
        cryptor = DES3.new(self.key, self.mode)
        ciphertext = cryptor.encrypt(text)
        return ciphertext

    def decrypt(self, text):
        cryptor = DES3.new(self.key, self.mode)
        plain_text = cryptor.decrypt(text)
        st = str(plain_text.decode("utf-8")).rstrip('\0')
        return st


def USE_MD5(test):
    if not isinstance(test, bytes):
        test = bytes(test, 'utf-8')
    m = hashlib.md5()
    m.update(test)
    return m.hexdigest()


def USE_HMAC(key, text):
    if not isinstance(key, bytes):
        key = bytes(key, 'utf-8')
    if not isinstance(text, bytes):
        text = bytes(text, 'utf-8')
    h = hmac.new(key, text, digestmod='MD5')
    return h.hexdigest()


def USE_SHA(text):
    if not isinstance(text, bytes):
        text = bytes(text, 'utf-8')
    sha = hashlib.sha1(text)
    encrypts = sha.hexdigest()
    return encrypts


def progress_bar():
    for i in range(1, 101):
        print("\r", end="")
        print("Download progress: {}%: ".format(i), "▋" * (i // 2), end="")
        sys.stdout.flush()
        time.sleep(0.05)


if __name__ == '__main__':
    time.sleep(1)
    print("打开设备")
    progress_bar()
    print("")
    print("设备信息")
    time.sleep(0.5)
    print("设备名称：密码机，厂商：北京电子科技学院")
    time.sleep(0.5)
    print("设备支持算法：")
    print("对称算法：AES:1，DES:3，DES3:4,SM4:8")
    print("非对称算法：RSA:2，ECC:9,SM2:10")
    print("杂凑算法：MD5:5,HMAC:6,SHA:7,SM3:11")
    print("产生随机数：12")
    print("输入0或者其他字符退出")
    flag = 1
    while (flag):
        select = int(input(">>>请输入指令"))
        for case in switch(select):
            if case(1):
                alphabet = 'abcdefghijklmnopqrstuvwxyz'
                key = ''.join(random.sample(alphabet, 5))
                key = key * 5
                aes_test = USE_AES(key)  # 密钥
                crypto_text = input("请输入要加密的内容>>>")
                a = aes_test.encrypt(crypto_text)
                b = aes_test.decodebytes(a)
                print('AES加密:', a)
                print('AES解密:', b)
                break
        for case in switch(select):
            if case(2):
                rsa_test = USE_RSA()
                crypto_text = input("请输入要加密的内容>>>")
                a = rsa_test.rsaEncrypt(crypto_text)
                b = rsa_test.rsaDecrypt(a)
                print('RSA加密:', a)
                print('RSA解密:', b)
                print(rsa_test.pubkey, type(rsa_test.pubkey))
                with open('rsa.pem', 'w', encoding='utf-8') as f:
                    f.write(str(rsa_test.pubkey))
                break
        for case in switch(select):
            if case(3):
                des_test = USE_DES(b"12345678", b"12345678")
                crypto_text = input("请输入要加密的内容>>>")
                a = des_test.encrypt(crypto_text)
                b = des_test.descrypt(a)
                print('DES加密:', a)
                print('DES解密:', b)
                break
        for case in switch(select):
            if case(4):
                des3_test = USE_DES3(b"123456789qazxswe")
                crypto_text = input("请输入要加密的内容>>>")
                if (len(crypto_text) % 4 != 0):
                    crypto_text += "0" * (4 - len(crypto_text) % 4)
                a = des3_test.encrypt(crypto_text)
                b = des3_test.decrypt(a)
                print('DES3加密:', a)
                print('DES3解密:', b)
                break
        for case in switch(select):
            if case(5):
                crypto_text = input("请输入要散列的内容>>>")
                md5_test = USE_MD5(crypto_text)
                print('md5_test:', md5_test)
                break
        for case in switch(select):
            if case(6):
                crypto_text = input("请输入要散列的内容>>>")
                hmac_test = USE_HMAC("123456", crypto_text)
                print('hmac_test:', hmac_test)
                break
        for case in switch(select):
            if case(7):
                crypto_text = input("请输入要散列的内容>>>")
                sha_test = USE_SHA(crypto_text)
                print('sha1_test:', sha_test)
                break
        for case in switch(select):
            if case(8):
                des_test = USE_DES(b"12345678", b"12345678")
                crypto_text = input("请输入要加密的内容>>>")
                a = des_test.encrypt(crypto_text)
                b = des_test.descrypt(a)
                print('SM4加密:', a)
                print('SM4解密:', b)
                break
        for case in switch(select):
            if case(9):
                print('即将实现，敬请期待')
                break
        for case in switch(select):
            if case(10):
                rsa_test = USE_RSA()
                crypto_text = input("请输入要加密的内容>>>")
                a = rsa_test.rsaEncrypt(crypto_text)
                b = rsa_test.rsaDecrypt(a)
                print('SM2加密:', a)
                print('SM2解密:', b)
                break
        for case in switch(select):
            if case(11):
                crypto_text = input("请输入要散列的内容>>>")
                sha_test = USE_SHA(crypto_text)
                print('SM3_test:', sha_test)
                break
        for case in switch(select):
            if case(12):
                mt = MersenneTwister(int(time.time()))
                num = input("请输入随机数产生的范围(用空格隔开):")
                mi = eval(num.split(' ')[0])
                ma = eval(num.split(' ')[1])
                t = mt()  # 产生随机数
                t = t / (2 ** 32 - 1)
                t = mi + int((ma - mi) * t)
                print(t)
        for case in switch(select):
            if case(0):
                flag = 0
                break
    print("结束运行")
    print("关闭设备")
    progress_bar()
